package com.electric.flinkcep.function;

import com.electric.flinkcep.model.LoadPredictionEvent;
import com.electric.flinkcep.model.PredictData;
import com.electric.flinkcep.model.TrueData;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;

/**
 * 数据匹配函数
 * 用于将真实数据和预测数据按时间戳进行匹配，生成LoadPredictionEvent
 */
public class DataMatchFunction extends CoProcessFunction<TrueData, PredictData, LoadPredictionEvent> {
    
    private static final Logger logger = LoggerFactory.getLogger(DataMatchFunction.class);
    
    // 存储等待匹配的真实数据
    private transient ValueState<Map<String, TrueData>> trueDataBuffer;
    
    // 存储等待匹配的预测数据
    private transient ValueState<Map<String, PredictData>> predictDataBuffer;
    
    // 时间容忍度（毫秒）- 允许5分钟的时间差
    private static final long TIME_TOLERANCE_MS = 5 * 60 * 1000;
    
    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        
        // 初始化状态
        ValueStateDescriptor<Map<String, TrueData>> trueDataDescriptor = 
            new ValueStateDescriptor<>("trueDataBuffer", 
                TypeInformation.of(new TypeHint<Map<String, TrueData>>() {}));
        trueDataBuffer = getRuntimeContext().getState(trueDataDescriptor);
        
        ValueStateDescriptor<Map<String, PredictData>> predictDataDescriptor = 
            new ValueStateDescriptor<>("predictDataBuffer", 
                TypeInformation.of(new TypeHint<Map<String, PredictData>>() {}));
        predictDataBuffer = getRuntimeContext().getState(predictDataDescriptor);
    }
    
    @Override
    public void processElement1(TrueData trueData, Context ctx, Collector<LoadPredictionEvent> out) throws Exception {
        logger.info("处理真实数据: {}", trueData);
        
        // 获取当前缓存的预测数据
        Map<String, PredictData> predictBuffer = predictDataBuffer.value();
        if (predictBuffer == null) {
            predictBuffer = new HashMap<>();
        }
        
        // 尝试找到匹配的预测数据
        PredictData matchedPredict = findMatchingPredictData(trueData, predictBuffer);
        
        if (matchedPredict != null) {
            // 找到匹配的预测数据，生成事件
            LoadPredictionEvent event = new LoadPredictionEvent(trueData, matchedPredict);
            logger.info("数据匹配成功: 真实负荷={}, 预测负荷={}, 误差={}%", 
                       trueData.getLoad(), matchedPredict.getPredictedLoad(), 
                       event.getRelativeErrorPercentage());
            out.collect(event);
            
            // 从缓存中移除已匹配的预测数据
            predictBuffer.remove(matchedPredict.getTimestamp());
            predictDataBuffer.update(predictBuffer);
        } else {
            // 没有找到匹配的预测数据，将真实数据放入缓存
            Map<String, TrueData> trueBuffer = trueDataBuffer.value();
            if (trueBuffer == null) {
                trueBuffer = new HashMap<>();
            }
            trueBuffer.put(trueData.getTimestamp(), trueData);
            trueDataBuffer.update(trueBuffer);
            
            logger.info("真实数据暂存，等待匹配的预测数据: {}", trueData.getTimestamp());
        }
        
        // 清理过期数据
        cleanupExpiredData(ctx.timestamp());
    }
    
    @Override
    public void processElement2(PredictData predictData, Context ctx, Collector<LoadPredictionEvent> out) throws Exception {
        logger.info("处理预测数据: {}", predictData);
        
        // 获取当前缓存的真实数据
        Map<String, TrueData> trueBuffer = trueDataBuffer.value();
        if (trueBuffer == null) {
            trueBuffer = new HashMap<>();
        }
        
        // 尝试找到匹配的真实数据
        TrueData matchedTrue = findMatchingTrueData(predictData, trueBuffer);
        
        if (matchedTrue != null) {
            // 找到匹配的真实数据，生成事件
            LoadPredictionEvent event = new LoadPredictionEvent(matchedTrue, predictData);
            logger.info("数据匹配成功: 真实负荷={}, 预测负荷={}, 误差={}%", 
                       matchedTrue.getLoad(), predictData.getPredictedLoad(), 
                       event.getRelativeErrorPercentage());
            out.collect(event);
            
            // 从缓存中移除已匹配的真实数据
            trueBuffer.remove(matchedTrue.getTimestamp());
            trueDataBuffer.update(trueBuffer);
        } else {
            // 没有找到匹配的真实数据，将预测数据放入缓存
            Map<String, PredictData> predictBuffer = predictDataBuffer.value();
            if (predictBuffer == null) {
                predictBuffer = new HashMap<>();
            }
            predictBuffer.put(predictData.getTimestamp(), predictData);
            predictDataBuffer.update(predictBuffer);
            
            logger.debug("预测数据暂存，等待匹配的真实数据: {}", predictData.getTimestamp());
        }
        
        // 清理过期数据
        cleanupExpiredData(ctx.timestamp());
    }
    
    /**
     * 查找匹配的预测数据
     */
    private PredictData findMatchingPredictData(TrueData trueData, Map<String, PredictData> predictBuffer) {
        long trueTimestamp = trueData.getTimestampMillis();
        
        for (PredictData predictData : predictBuffer.values()) {
            long predictTimestamp = predictData.getTimestampMillis();
            if (Math.abs(trueTimestamp - predictTimestamp) <= TIME_TOLERANCE_MS) {
                return predictData;
            }
        }
        return null;
    }
    
    /**
     * 查找匹配的真实数据
     */
    private TrueData findMatchingTrueData(PredictData predictData, Map<String, TrueData> trueBuffer) {
        long predictTimestamp = predictData.getTimestampMillis();
        
        for (TrueData trueData : trueBuffer.values()) {
            long trueTimestamp = trueData.getTimestampMillis();
            if (Math.abs(trueTimestamp - predictTimestamp) <= TIME_TOLERANCE_MS) {
                return trueData;
            }
        }
        return null;
    }
    
    /**
     * 清理过期数据
     */
    private void cleanupExpiredData(long currentTimestamp) throws Exception {
        // 清理过期的真实数据
        Map<String, TrueData> trueBuffer = trueDataBuffer.value();
        if (trueBuffer != null) {
            trueBuffer.entrySet().removeIf(entry -> {
                long dataTimestamp = entry.getValue().getTimestampMillis();
                return (currentTimestamp - dataTimestamp) > TIME_TOLERANCE_MS * 2;
            });
            trueDataBuffer.update(trueBuffer);
        }
        
        // 清理过期的预测数据
        Map<String, PredictData> predictBuffer = predictDataBuffer.value();
        if (predictBuffer != null) {
            predictBuffer.entrySet().removeIf(entry -> {
                long dataTimestamp = entry.getValue().getTimestampMillis();
                return (currentTimestamp - dataTimestamp) > TIME_TOLERANCE_MS * 2;
            });
            predictDataBuffer.update(predictBuffer);
        }
    }
}